#
# QsarDB class model
#

library("pmml")
library("XML")

setClass("qdb", representation(dir = "character"))

setClass("container", representation(id = "character", name = "character", labels = "vector"), prototype = list(id = NA_character_, name = NA_character_, labels = vector()))
setClass("parameter", contains = "container")
setClass("compound", representation(cas = "character", inChI = "character"), prototype = list(cas = NA_character_, inChI = NA_character_), contains = "container")
setClass("descriptor", contains = "parameter")
setClass("property", representation(endpoint = "character", species = "character"), prototype = list(endpoint = NA_character_, species = NA_character_), contains = "parameter")
setClass("prediction", contains = "parameter")
setClass("model", contains = "container")

setClass("containerRegistry", representation(), contains = "list")
setClass("parameterRegistry", contains = "containerRegistry")
setClass("compoundRegistry", contains = "containerRegistry")
setClass("descriptorRegistry", contains = "parameterRegistry")
setClass("propertyRegistry", contains = "parameterRegistry")
setClass("predictionRegistry", contains = "parameterRegistry")
setClass("modelRegistry", contains = "containerRegistry")

loadContainerRegistry = function(path, containerRegistryClass, containerClass){
	doc = xmlTreeParse(path, useInternalNodes = TRUE)

	capitalize = function(x){
		return (paste(toupper(substring(x, 1, 1)), substring(x, 2), sep = ""))
	}

	containerNodes = getNodeSet(doc, paste("", capitalize(containerRegistryClass), capitalize(containerClass), sep = "/qdb:"), namespaces = c(qdb = "http://www.qsardb.org/QDB"))

	containerRegistry = new(containerRegistryClass)

	stringValue = function(parent, name){
		return (xmlValue(parent[[name]]))
	}

	stringListValue = function(parent, name){
		return (unlist(strsplit(xmlValue(parent[[name]]), split = " ")))
	}

	containerList = list()

	for(containerNode in containerNodes){
		container = new(containerClass)

		# string-valued attributes
		container@id = stringValue(containerNode, "Id")
		container@name = stringValue(containerNode, "Name")

		# string list-values attributes
		container@labels = stringListValue(containerNode, "Labels")

		containerList[container@id] = container
	}

	containerRegistry = c(containerList)

	return (containerRegistry)
}

loadCompoundRegistry = function(qdb){
	return (loadContainerRegistry(paste(qdb@dir, "compounds/compounds.xml", sep = "/"), "compoundRegistry", "compound"))
}

loadPropertyRegistry = function(qdb){
	return (loadContainerRegistry(paste(qdb@dir, "properties/properties.xml", sep = "/"), "propertyRegistry", "property"))
}

loadDescriptorRegistry = function(qdb){
	return (loadContainerRegistry(paste(qdb@dir, "descriptors/descriptors.xml", sep = "/"), "descriptorRegistry", "descriptor"))
}

loadModelRegistry = function(qdb){
	return (loadContainerRegistry(paste(qdb@dir, "models/models.xml", sep = "/"), "modelRegistry", "model"))
}

loadPredictionRegistry = function(qdb){
	return (loadContainerRegistry(paste(qdb@dir, "predictions/predictions.xml", sep = "/"), "predictionRegistry", "prediction"))
}

selectContainers = function(containerRegistry, filter){
	containerList = list()

	for(container in containerRegistry){
	
		if(filter(container)){
			containerList[container@id] = container
		}
	}
	
	return (containerList)
}

selectContainersById = function(containerRegistry, ids){
	containsId = function(container){
		return (container@id %in% ids)
	}

	return (selectContainers(containerRegistry, containsId))
}

cargoPath = function(qdb, container, cargo){
	type = ""

	if(class(container) == "compound"){
		type = "compounds"
	} else

	if(class(container) == "property"){
		type = "properties"
	} else

	if(class(container) == "descriptor"){
		type = "descriptors"
	} else

	if(class(container) == "model"){
		type = "models"
	} else

	if(class(container) == "prediction"){
		type = "predictions"
	} else

	{
		stop(class(container))
	}

	return (paste(qdb@dir, type, container@id, cargo, sep = "/"))
}

loadParameterValues = function(qdb, parameter){
	path = cargoPath(qdb, parameter, "values")

	return (read.table(file = path, header = TRUE, sep = "\t", na.strings = "N/A", colClasses = c("character", NA)));
}

storeParameterValues = function(qdb, parameter, values){
	path = cargoPath(qdb, parameter, "values")

	write.table(values, file = path, col.names = c("Compound Id", parameter@name), row.names = FALSE, sep = "\t", quote = FALSE)
}

escapeParameterId = function(id){
	return (gsub("[ ,-]", "_", x = id));
}

loadModelData = function(qdb, property, descriptors){
	parameters = c(list(property), descriptors)

	return (loadParameterData(qdb, parameters))
}

loadParameterData = function(qdb, parameters){

	loadColumn = function(qdb, parameter){
		result = loadParameterValues(qdb, parameter)
		names(result) = c("Id", escapeParameterId(parameter@id))

		return (result)
	}

	data = NULL

	for(parameter in parameters){
	
		if(is.null(data)){
			data = loadColumn(qdb, parameter)
		} else 
		
		{
			data = merge(data, loadColumn(qdb, parameter), by = "Id")
		}
	}

	return (data)
}

storeModelPMML = function(qdb, model, object){
	path = cargoPath(qdb, model, "pmml")

	saveXML(file = path, pmml(object))
}

loadModelRDS = function(qdb, model){
	path = cargoPath(qdb, model, "rds")

	return (readRDS(file = path))
}

storeModelRDS = function(qdb, model, object){
	path = cargoPath(qdb, model, "rds")

	saveRDS(file = path, object)
}